from core.function import Function
from utils.common import logsumexp
import utils

class LogSoftmax(Function):
    def __init__(self, axis=1):
        """
        实现 Log Softmax 操作的函数。

        Args:
            axis (int): 沿着哪个轴计算 Softmax。默认是 1。
        """
        self.axis = axis

    def forward(self, x):
        """
        计算 Log Softmax 操作的前向传播。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: Log Softmax 操作的输出变量。
        """
        log_z = logsumexp(x, self.axis)
        y = x - log_z
        return y

    def backward(self, gy):
        """
        计算 Log Softmax 操作的反向传播。

        Args:
            gy (Variable): 输出的梯度。

        Returns:
            Variable: 输入的梯度。
        """
        y = self.outputs[0]()
        gx = gy - utils.functions_collect.exp(y) * gy.sum(axis=self.axis, keepdims=True)
        return gx

    def log_softmax(self, x):
        """
        对输入变量应用 Log Softmax 操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: Log Softmax 操作的输出变量。
        """
        return self(x)
